#pragma once
#include <memory>
#include <vector>
#include <queue>
#include <functional>
#include <cmath>
#include <limits>
#include <algorithm>

template <typename T, typename V>
class KDTree3D
{
private:
    using QueueType = std::priority_queue<std::pair<V, T>,
                                          std::vector<std::pair<V, T>>,
                                          std::function<bool(const std::pair<V, T> &,
                                                             const std::pair<V, T> &)>>;
    struct Node
    {
        T point;
        std::shared_ptr<Node> left;
        std::shared_ptr<Node> right;
        Node(T p) : point(p), left(nullptr), right(nullptr) {}
    };

    int max_depth;
    std::shared_ptr<Node> root;
    std::shared_ptr<Node> buildBalancedTree(std::vector<T> &points, int depth)
    {
        if (depth > max_depth)
            max_depth = depth;
        if (points.empty())
            return nullptr;

        int axis = depth % 3;
        std::sort(points.begin(), points.end(), [this, axis](const T &a, const T &b)
                  { return targetCoord(a, axis) < targetCoord(b, axis); });

        size_t mid = points.size() / 2;
        auto node = std::make_shared<Node>(points[mid]);

        std::vector<T> leftPoints(points.begin(), points.begin() + mid);
        std::vector<T> rightPoints(points.begin() + mid + 1, points.end());
        node->left = buildBalancedTree(leftPoints, depth + 1);
        node->right = buildBalancedTree(rightPoints, depth + 1);

        return node;
    }
    void nearestNeighborRec(std::shared_ptr<Node> node, T target, int depth,
                            T &best, V &bestDist)
    {
        if (!node)
            return;

        V dist = distance(node->point, target);
        if (dist < bestDist)
        {
            bestDist = dist;
            best = node->point;
        }

        int axis = depth % 3;
        bool goLeft = (axis == 0 && target.x < node->point.x) ||
                      (axis == 1 && target.y < node->point.y) ||
                      (axis == 2 && target.z < node->point.z);

        if (goLeft)
        {
            nearestNeighborRec(node->left, target, depth + 1, best, bestDist);
            if (abs(targetCoord(target, axis) - targetCoord(node->point, axis)) < bestDist)
            {
                nearestNeighborRec(node->right, target, depth + 1, best, bestDist);
            }
        }
        else
        {
            nearestNeighborRec(node->right, target, depth + 1, best, bestDist);
            if (abs(targetCoord(target, axis) - targetCoord(node->point, axis)) < bestDist)
            {
                nearestNeighborRec(node->left, target, depth + 1, best, bestDist);
            }
        }
    }
    void KnearestNeighborsRec(std::shared_ptr<Node> node, T target, int depth,
                              int k, QueueType &kNearest)
    {
        if (!node)
            return;

        V dist = distance(node->point, target);

        if (kNearest.size() < k)
        {
            kNearest.emplace(dist, node->point);
        }
        else if (dist < kNearest.top().first)
        {
            kNearest.pop();
            kNearest.emplace(dist, node->point);
        }

        int axis = depth % 3;
        bool goLeft = (axis == 0 && target.x < node->point.x) ||
                      (axis == 1 && target.y < node->point.y) ||
                      (axis == 2 && target.z < node->point.z);

        if (goLeft)
        {
            KnearestNeighborsRec(node->left, target, depth + 1, k, kNearest);
            if (std::abs(targetCoord(target, axis) - targetCoord(node->point, axis)) < kNearest.top().first)
            {
                KnearestNeighborsRec(node->right, target, depth + 1, k, kNearest);
            }
        }
        else
        {
            KnearestNeighborsRec(node->right, target, depth + 1, k, kNearest);
            if (std::abs(targetCoord(target, axis) - targetCoord(node->point, axis)) < kNearest.top().first)
            {
                KnearestNeighborsRec(node->left, target, depth + 1, k, kNearest);
            }
        }
    }
    void printTreeRec(std::shared_ptr<Node> node, int depth, std::string prefix, bool isLeft) const
    {
        if (node == nullptr)
            return;

        std::cout << prefix;

        std::cout << (isLeft ? "|-- " : "\\-- ");

        std::cout << "Node: " << node.get()
                  << " (Point: " << "(" << node->point.x << ", " << node->point.y << ", " << node->point.z << "))"
                  << ", Left: " << node->left.get()
                  << ", Right: " << node->right.get() << "\n";

        printTreeRec(node->left, depth + 1, prefix + (isLeft ? "|   " : "    "), true);
        printTreeRec(node->right, depth + 1, prefix + (isLeft ? "|   " : "    "), false);
    }
    void Interpolation(T &target, QueueType kNearest)
    {
        V weight_S1 = 0.0, weight_sum = 0.0;
        while (!kNearest.empty())
        {
            V dist = kNearest.top().first;
            T tmp = kNearest.top().second;
            
            if(dist == 0.0)
            {
                target.s1 = tmp.s1;
                return;
            }
            V inv_dist = (V)(1.0 / dist);

            weight_sum += inv_dist;
            weight_S1 += tmp.s1 * inv_dist;
            kNearest.pop();
        }

        target.s1 = weight_S1 / weight_sum;
    }
    V distance(const T &a, const T &b) const
    {
        return std::sqrt((a.x - b.x) * (a.x - b.x) +
                         (a.y - b.y) * (a.y - b.y) +
                         (a.z - b.z) * (a.z - b.z));
    }
    V targetCoord(const T &point, int axis) const
    {
        return (axis == 0) ? point.x : (axis == 1) ? point.y
                                                   : point.z;
    }
    static bool cmp(const std::pair<V, T> &a, const std::pair<V, T> &b)
    {
        return a.first < b.first;
    }

public:
    KDTree3D(const std::vector<T> &points)
    {
        max_depth = 0;
        std::vector<T> pointsCopy = points;
        root = buildBalancedTree(pointsCopy, 0);
        std::cout << "mytree max depth is " << max_depth << std::endl;
    }
    T nearestNeighbor(const T &target)
    {
        T best;
        best.x = 0;
        best.y = 0;
        best.z = 0;
        best.s1 = 0;
        V bestDist = std::numeric_limits<V>::max();
        nearestNeighborRec(root, target, 0, best, bestDist);
        return best;
    }
    std::vector<T> KnearestNeighbors(const T &target, int k)
    {
        QueueType kNearest(cmp);

        KnearestNeighborsRec(root, target, 0, k, kNearest);

        std::vector<T> result;
        while (!kNearest.empty())
        {
            result.push_back(kNearest.top().second);
            kNearest.pop();
        }
        std::reverse(result.begin(), result.end());
        return result;
    }
    void K_Interpolation(T &target, int k)
    {
        QueueType kNearest(cmp);

        KnearestNeighborsRec(root, target, 0, k, kNearest);

        Interpolation(target, kNearest);
    }
    void printTree() const
    {
        printTreeRec(root, 0, "", true);
    }
};